from axes_divider import Size, LocatableAxes, make_axes_locatable, \
     Divider, SubplotDivider

import matplotlib.pyplot as plt

def test1():
    fig1 = plt.figure(1)
    fig1.clf()
    pos = [0.1, 0.1, 0.8, 0.8] 

    axes1 = LocatableAxes(fig1, pos)
    axes2 = LocatableAxes(fig1, pos)
    axes3 = LocatableAxes(fig1, pos)

    aspectx, aspecty = 1., 1.
    
    h = [Size.Bound(axes1.get_xlim, aspectx),
         Size.Fixed(0.1),
         Size.Fraction(Size.Bound(axes1.get_xlim, aspectx), .5),
         ]
    
    v = [Size.Bound(axes1.get_ylim, aspecty),
         Size.Fixed(0.1),
         Size.Bound(axes3.get_ylim, aspecty),]
    
    axes_divider = Divider(fig1, pos,
                           horizontal=h,
                           vertical=v, aspect=True)

    axes_locator1 = axes_divider.new_locator(nx=0, ny=0)
    axes_locator2 = axes_divider.new_locator(nx=2, ny=0)
    axes_locator3 = axes_divider.new_locator(nx=0, ny=2, nx1=3)

    axes1.set_axes_locator(axes_locator1)
    axes2.set_axes_locator(axes_locator2)
    axes3.set_axes_locator(axes_locator3)

    fig1.add_axes(axes1)
    fig1.add_axes(axes2)
    fig1.add_axes(axes3)

    axes1.set_ylim(0., 2.)
    plt.draw()

    

def test2():
    fig1 = plt.figure(1)
    fig1.clf()
    pos = [0.1, 0.1, 0.8, 0.8] 

    axes1 = LocatableAxes(fig1, pos)
    lab_ax = LocatableAxes(fig1, pos, frameon=False)
    lab_ax.xaxis.set_visible(False)
    lab_ax.yaxis.set_visible(False)

    lab_ax2 = LocatableAxes(fig1, pos, frameon=False)
    lab_ax2.xaxis.set_visible(False)
    lab_ax2.yaxis.set_visible(False)

    p1, = axes1.plot([1,3,2])
    p2, = axes1.plot([1,2,2.5])

    lgnd = lab_ax.legend([p1, p2, p1, p2], ["test1","test2", "t3", "t4"],
                         borderaxespad=0., ncol=4, mode="expand")
    #lgnd.get_frame().set_visible(False)

    lgnd2 = lab_ax2.legend([p1, p2, p1, p2], ["test1","test2", "t3", "t4"],
                           borderaxespad=0., ncol=1, loc=2)
    #lgnd2.get_frame().set_visible(False)

    aspectx, aspecty = 1., 1.
    
    h = [Size.Scalable(1),
         Size.MaxExtent([lgnd2._legend_box], "width")
         ]
    
    v = [Size.Scalable(1),
         #MaxExtentSize(lgnd.get_texts(), "height")
         Size.MaxExtent([lgnd._legend_box], "height")
         ]
    
    axes_divider = Divider(fig1, pos,
                           horizontal=h,
                           vertical=v, aspect=True)

    axes_locator1 = axes_divider.new_locator(nx=0, ny=0)
    axes_locator2 = axes_divider.new_locator(nx=0, ny=1)
    axes_locator3 = axes_divider.new_locator(nx=1, ny=0)

    axes1.set_axes_locator(axes_locator1)
    lab_ax.set_axes_locator(axes_locator2)
    lab_ax2.set_axes_locator(axes_locator3)

    fig1.add_axes(axes1)
    fig1.add_axes(lab_ax)
    fig1.add_axes(lab_ax2)

    #axes1.set_ylim(0., 2.)
    plt.draw()


def get_image():
    # prepare image
    delta = 0.5

    extent = (-3,4,-4,3)
    import numpy as np
    x = np.arange(-3.0, 4.001, delta)
    y = np.arange(-4.0, 3.001, delta)
    X, Y = np.meshgrid(x, y)
    import matplotlib.mlab as mlab
    Z1 = mlab.bivariate_normal(X, Y, 1.0, 1.0, 0.0, 0.0)
    Z2 = mlab.bivariate_normal(X, Y, 1.5, 0.5, 1, 1)
    Z = (Z1 - Z2) * 10

    return Z, extent

def test3():
    fig1 = plt.figure(1)
    fig1.clf()


    axes1 = fig1.add_subplot(2, 1, 1)


    Z, extent = get_image()
    im = axes1.imshow(Z, extent=extent, interpolation="nearest")

    axes2 = fig1.add_subplot(2, 1, 2)
    p1, = axes2.plot([1,3,2])
    p2, = axes2.plot([1,2,2.5])
    p3, = axes2.plot([2,1,0.5])

    ### 
    axes_divider = make_axes_locatable(axes1)

    axes_cb = axes_divider.new_horizontal(size="5%", pad="5%")
    fig1.add_axes(axes_cb)
    
    axes_cb.yaxis.set_ticks_position("right")
    plt.colorbar(im, cax=axes_cb)


    axes_divider = make_axes_locatable(axes2)

    #lgnd = axes2.legend([p1, p2, p3], ["test1", "test2", "test3"],
    #                    borderaxespad=0., ncol=3, mode="expand", loc=(0., 1.05))

    mysize = Size.MaxExtent([], "height")
    axes_legend = axes_divider.new_vertical(size=mysize, pad=0.05)
    axes_legend.xaxis.set_visible(False)
    axes_legend.yaxis.set_visible(False)
    
    lgnd = axes_legend.legend([p1, p2, p3], ["test1", "test2", "test3"],
                              borderaxespad=0., ncol=3, mode="expand")
    mysize.add_artist(lgnd._legend_box)

    fig1.add_axes(axes_legend)
    

    plt.draw()




def test4():
    import matplotlib.pyplot as plt

    fig1 = plt.figure(1)
    fig1.clf()


    divider = SubplotDivider(fig1, 1, 1, 1)
    ax = LocatableAxes(fig1, divider.get_position())
    ax_cb = LocatableAxes(fig1, divider.get_position())

    p1, = ax.plot([1,3,2])
    p2, = ax.plot([1,2,2.5])
    p3, = ax.plot([2,1,0.5])

    ### 
    h = [Size.AxesX(ax),
         Size.Fixed(0.1),
         Size.Fixed(0.5),
         ]

    v = [Size.AxesX(ax),
         ]

    divider._horizontal = h
    divider._vertical = v

    ax.set_axes_locator(divider.new_locator(nx=0, ny=0))
    ax_cb.set_axes_locator(divider.new_locator(nx=2, ny=0))

    fig1.add_axes(ax)
    fig1.add_axes(ax_cb)
    
    ax_cb.yaxis.set_ticks_position("right")

    plt.draw()
    #plt.colorbar(im, cax=ax_cb)
    
